#pragma warning(disable:4661)

#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include "helper_math.h"

#include "CudaIcemVolume.h"
#include "IcemVolume.h"
#include "TalairachVolume.h"
#include "IcemStack.h"


#include <cstdint>
#include <stdio.h>
#include <vector>
#include <ctime>
#include <fstream>


//template<class T>

texture<unsigned char, cudaTextureType3D,cudaReadModeElementType> volumeTexture;
texture<unsigned short, cudaTextureType3D,cudaReadModeElementType> volumeTextureTalairach;


template <class T>
void IcemVolume<T>::GetRenderedVolume(unsigned char* v, double minCoord_in[3], size_t size[3] ,bool interp){
	double pos[3] = {minCoord_in[0],0,0};
	SlicePreparation(pos, 2, minCoord_in);
	int ret = cudaMalloc(&cudaMat,16*sizeof(float));
	unsigned char* cudavolume_out;
	ret = cudaMalloc(&cudavolume_out, size[0]*size[1]*size[2]/8);
	ret = cudaMemcpy(cudaMat,sliceTransformation,16*sizeof(float),cudaMemcpyHostToDevice);

	cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<T>();

	ret = cudaBindTextureToArray(&volumeTexture,(cudaArray*)cudaVolume,&channelDesc);
	volumeTexture.addressMode[0] = cudaAddressModeClamp;
	volumeTexture.addressMode[1] = cudaAddressModeClamp;
	volumeTexture.addressMode[2] = cudaAddressModeClamp;
	volumeTexture.filterMode = cudaFilterModePoint;
	volumeTexture.normalized = false;

	//CudaRenderedVolume(float contrast, float brightness, float step, int sizeI, int sizeJ, int sizeK, int sizeVol1, 
	//int sizeVol2, int sizeVol3, unsigned char* volume_out, float* mat, bool interpolate){
	CudaRenderedVolume<T><<<NBLOCKS,NTHREADS>>> (contrast, brightness, 1.f, size[2]/2, size[1]/2, size[0]/2,shape[0]*shape[1],shape[0],
								shape[0]*shape[1]*shape[2],cudavolume_out,cudaMat,interp);
	ret = cudaDeviceSynchronize();
	cudaMemcpy(v,cudavolume_out, size[0]*size[1]*size[2]/8,cudaMemcpyDeviceToHost);
	cudaFree(cudavolume_out);
	cudaFree(cudaMat);
}


short3* TalairachVolume::GetSlice(double pos[3], int dir, double minCoord_in[3], bool firstSlice, bool interp){
	SlicePreparation(pos, dir, minCoord_in);
	int ret = cudaMalloc(&cudaMat,16*sizeof(float));
	ret = cudaMalloc(&cu_invT,16*sizeof(float));
	ret = cudaMemcpy(cu_invT,invT,16*sizeof(float),cudaMemcpyHostToDevice);
	ret = cudaMemcpy(cudaMat,sliceTransformation,16*sizeof(float),cudaMemcpyHostToDevice);
	ret = cudaMalloc(&cuTalPoints,12*sizeof(float3));
	ret = cudaMemcpy(cuTalPoints,talairachPoints,12*sizeof(float3),cudaMemcpyHostToDevice);

	cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<uint16_t>();

	ret = cudaBindTextureToArray(&volumeTextureTalairach,(cudaArray*)cudaVolume,&channelDesc);
	volumeTextureTalairach.addressMode[0] = cudaAddressModeClamp;
	volumeTextureTalairach.addressMode[1] = cudaAddressModeClamp;
	volumeTextureTalairach.addressMode[2] = cudaAddressModeClamp;
	volumeTextureTalairach.filterMode = cudaFilterModePoint;
	volumeTextureTalairach.normalized = false;
	CudaTalairachSlice<<<NBLOCKS,NTHREADS>>> (contrast,brightness,step, sliceSizeI[dir], sliceSizeJ[dir],shape[0]*shape[1],shape[0],
								shape[0]*shape[1]*shape[2],cudaImage[dir],cudaMat, colormap,cuTalPoints, dir==2, cu_invT);
	ret = cudaDeviceSynchronize();
	cudaFree(cudaMat);
	cudaFree(cu_invT);
	cudaFree(cuTalPoints);
	return cudaImage[dir];
}

template <class T>
short3* IcemVolume<T>::GetSlice(double pos[3], int dir, double minCoord_in[3], bool firstSlice, bool interp){
	SlicePreparation(pos, dir, minCoord_in);
	int ret = cudaMalloc(&cudaMat,16*sizeof(float));
	ret = cudaMemcpy(cudaMat,sliceTransformation,16*sizeof(float),cudaMemcpyHostToDevice);

	cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<T>();

	ret = cudaBindTextureToArray(&volumeTexture,(cudaArray*)cudaVolume,&channelDesc);
	volumeTexture.addressMode[0] = cudaAddressModeClamp;
	volumeTexture.addressMode[1] = cudaAddressModeClamp;
	volumeTexture.addressMode[2] = cudaAddressModeClamp;
	volumeTexture.filterMode = cudaFilterModePoint;
	volumeTexture.normalized = false;
	CudaSlice<T><<<NBLOCKS,NTHREADS>>> (contrast, brightness, step, sliceSizeI[dir], sliceSizeJ[dir],shape[0]*shape[1],shape[0],
								shape[0]*shape[1]*shape[2],cudaImage[dir],cudaMat, colormap, dir==2, firstSlice,interp, filterLimits);
	ret = cudaDeviceSynchronize();
	cudaFree(cudaMat);
	return cudaImage[dir];
}
	
__device__ void mult_backup(const unsigned char& alpha, uchar3& in) {
	unsigned short s;
	s  = unsigned short(in.x)*unsigned short(alpha);
	in.x = s >> 8;
	s  = in.y*alpha;
	in.y = s >> 8;
	s  = in.z*alpha;
	in.z = s >> 8;
};

__device__ uchar3 mult(float alpha, uchar3 in) {
	in.x = alpha * in.x;
	in.y = alpha * in.y;
	in.z = alpha * in.z;
	return in;
};


__global__ void CudaStack(int dim, uchar3* cuOut, short3** cuImageStack,float* cuAlpha, int sizeI, int sizeJ, bool transpose){
	//int i=threadIdx.x;
	__shared__ float salpha[20];
	if (threadIdx.x < dim)
		salpha[threadIdx.x] = cuAlpha[threadIdx.x];
	__syncthreads();
	for (int i=threadIdx.x;i<sizeI;i+=NTHREADS) {
		for (int j=blockIdx.x;j<sizeJ;j+=NBLOCKS) {
			uchar3 out = make_uchar3(0,0,0);
			float alpha_b = 1;
			for (int layer=0;layer<dim;++layer){
				float alpha_a = salpha[layer];
				float alpha0 = (alpha_a + alpha_b * (1-alpha_a));
				short3 sx = transpose?cuImageStack[layer][sizeI * j + i]:cuImageStack[layer][sizeJ * i + j];
				uchar3 x;
				if (sx.x<0) { //x<0
					alpha_a=0;
					alpha0 = alpha_b; //alpha_a + alpha_b * (1-alpha_a);
					x = make_uchar3(0,0,0);
				} else {
					x.x = sx.x;
					x.y = sx.y;
					x.z = sx.z;
				}
				float a1 = alpha_b*(1-alpha_a)/alpha0;
				float a2 = alpha_a/alpha0;
				out = mult(a1,out);
				x = mult(a2,x);
				out.x = out.x + x.x;
				out.y = out.y + x.y;
				out.z = out.z + x.z;
				alpha_b = alpha0;
			}
			if (!transpose) {
				cuOut[sizeJ * i + j] = out;
            } else {
				cuOut[sizeI * j + i] = out;
            }
			//}
		}
	}
}

	__global__ void CudaTalairachSlice(float contrast, float brightness, float step,int sizeI, int sizeJ, int sizeVol1, 
		int sizeVol2, int sizeVol3, short3* im, float* mat, unsigned char * colormap, float3* cuTalPoints, bool transpose, float* cu_invT){
		//__shared__ float scolormap[3*2000];
		__shared__ float smat[16];
		__shared__ float3 sCuTalPoints[12];
		
		if (threadIdx.x<16)
			smat[threadIdx.x] = mat[threadIdx.x];
		if (threadIdx.x<12)
			sCuTalPoints[threadIdx.x] = cuTalPoints[threadIdx.x];
		//for (int k=threadIdx.x;k<3*2000;k+=NTHREADS) {
		//		scolormap[k] = colormap[k];
		//}
		__syncthreads();
		int N_i = sizeI/NTHREADS + 1;
		int N_j = sizeJ/NBLOCKS + 1;
		for (int i=threadIdx.x*N_i;i<min(sizeI,(threadIdx.x+1)*N_i);++i) {
			float M1[4];
			float M[4];
			M1[0] = i*step;
			M1[2] = 0;
			M1[3] = 1;
			for (int j= blockIdx.x*N_j;j<min(sizeJ,(blockIdx.x+1)*N_j);++j) {
				M1[1] = j* step;
				
				float3 in, out;
				in = make_float3(M1[0],M1[1],M1[2]);
				out  = in;//Talairach(in,sCuTalPoints);
				M1[0] = out.x;
				M1[1] = out.y;
				M1[2] = out.z;
				for (int k=0;k<4;++k) {
					M[k] = 0;
					for (int l=0;l<4;++l) {
						M[k] += smat[k*4 + l] * M1[l];
					}
				}
				in = make_float3(M[0],M[1],M[2]);
				out  = Talairach(in,sCuTalPoints);
				M[0] = out.x;
				M[1] = out.y;
				M[2] = out.z;
				float M2[4];
				for (int k=0;k<4;++k) {
					M2[k] = 0;
					for (int l=0;l<4;++l) {
						M2[k] += cu_invT[k*4 + l] * M[l];
					}
				}
				int x,y,z;
				//round to nearest neighbor
				x = (int) floor( M2[0] +0.5 );
				y = (int) floor( M2[1] +0.5 );
				z = (int) floor( M2[2]+0.5 );
				int sizeX = sizeVol2;
				int sizeY = sizeVol1/sizeVol2;
				int sizeZ = sizeVol3/sizeVol1;
				//int index = sizeVol1 * z + sizeVol2 * y + x;
				float value = 0;
				if (x>=0 && x<sizeX && y>=0 && y<sizeY && z>=0 && z<sizeZ) {
					value = (float)tex3D(volumeTextureTalairach,x,y,z);		
				}
				bool isNull = (value==0);
				unsigned short value2(value);
				short r,g,b;				
				r = isNull?-1:colormap[value2];
				g = isNull?-1:colormap[2000 + value2];
				b = isNull?-1:colormap[4000 + value2];
				short3 rgb = make_short3(r,g,b);
				if (transpose) {
					im[sizeI*j+i] = rgb;
				}
				else {
					im[sizeJ*i+j] = rgb;
				}				
			}
		}		
	}


template <class T> 
__global__ void CudaSlice(float contrast, float brightness, float step,int sizeI, int sizeJ, int sizeVol1, 
	int sizeVol2, int sizeVol3, short3* im, float* mat, unsigned char* colormap, bool transpose, bool firstSlice,
    bool interpolate, float2 filterLimits){
	//__shared__ float scolormap[3*256];
	__shared__ float smat[16];
	if (threadIdx.x<16)
		smat[threadIdx.x] = mat[threadIdx.x];
	//for (int k=threadIdx.x;k<3*256;k+=NTHREADS) {
	//		scolormap[k] = colormap[k];
	//}
	__syncthreads();
	int N_i = sizeI/NTHREADS + 1;
	int N_j = sizeJ/NBLOCKS + 1;
	for (int i=threadIdx.x*N_i;i<min(sizeI,(threadIdx.x+1)*N_i);++i) {
		float M1[4];
		float M[4];
		M1[0] = i*step;
		M1[2] = 0;
		M1[3] = 1;
		for (int j= blockIdx.x*N_j;j<min(sizeJ,(blockIdx.x+1)*N_j);++j) {
			M1[1] = j* step;
			for (int k=0;k<3;++k) {
				M[k] = 0;
				for (int l=0;l<4;++l) {
					M[k] += smat[k*4 + l] * M1[l];
				}
			}
			int x,y,z;
			if (!interpolate) {//round to nearest neighbor
				x = (int) floor( M[0]+0.5 );
				y = (int) floor( M[1]+0.5 );
				z = (int) floor( M[2]+0.5 );
			}
			else {
				x = (int) floor(M[0]);
				y = (int) floor(M[1]);
				z = (int) floor(M[2]);
			}
			int sizeX = sizeVol2;
			int sizeY = sizeVol1/sizeVol2;
			int sizeZ = sizeVol3/sizeVol1;
			//int index = sizeVol1 * z + sizeVol2 * y + x;
			float value = 0;
			bool isNull = !firstSlice;
			if (interpolate && x>=0 && x<sizeX-1 && y>=0 && y<sizeY-1 && z>=0 && z<sizeZ-1) {

			//trilinear interpolation, see http://en.wikipedia.org/wiki/Trilinear_interpolation
				float dx,dy,dz;
				dx = (M[0] - x);//:step;
				dy = (M[1] - y);///step;
				dz = (M[2] - z);///step;

				float c000 = (float)tex3D(volumeTexture,x,y,z);					
				float c001 = (float)tex3D(volumeTexture,x,y,z+1);			
				float c010 = (float)tex3D(volumeTexture,x,y+1,z);
				float c011 = (float)tex3D(volumeTexture,x,y+1,z+1);
				float c100 = (float)tex3D(volumeTexture,x+1,y,z);
				float c101 = (float)tex3D(volumeTexture,x+1,y,z+1);
				float c110 = (float)tex3D(volumeTexture,x+1,y+1,z);
				float c111 = (float)tex3D(volumeTexture,x+1,y+1,z+1);
				float i1 = c000 * (1-dz) + c001 * dz;
				float i2 = c010 * (1-dz) + c011 * dz;
				float j1 = c100 * (1-dz) + c101 * dz;
				float j2 = c110 * (1-dz) + c111 * dz;
				float w1 = i1   * (1-dy) + i2   * dy;
				float w2 = j1   * (1-dy) + j2   * dy;
				value    = w1   * (1-dx) + w2   * dx;
			}
			else if (x>=0 && x<sizeX && y>=0 && y<sizeY && z>=0 && z<sizeZ) {
				value = (float)tex3D(volumeTexture,x,y,z);		
			}
			isNull &= (value==0 || (value>filterLimits.x && value<filterLimits.y) );
			if (!isNull) {
				if (brightness < 0)
					value *= 1 + brightness;
				else
					value += (255-value) * brightness;
				value = (value - 128) * tanf((contrast+1) * PI / 4.) + 128;
			}
			unsigned char value2(value);
			short r,g,b;				
			r = isNull?-1:colormap[value2];
			g = isNull?-1:colormap[256 + value2];
			b = isNull?-1:colormap[512 + value2];
			short3 rgb = make_short3(r,g,b);
			if (transpose) {
				im[sizeI*j+i] = rgb;
			}
			else {
				im[sizeJ*i+j] = rgb;
			}
		}
	}		
}

template <class T> 
__global__ void CudaRenderedVolume(float contrast, float brightness, float step, int sizeI, int sizeJ, int sizeK, int sizeVol1, 
	int sizeVol2, int sizeVol3, unsigned char* volume_out, float* mat, bool interpolate){
	__shared__ float smat[16];
	if (threadIdx.x<16)
		smat[threadIdx.x] = mat[threadIdx.x];
	__syncthreads();
	int N_i = sizeI/NTHREADS + 1;
	int N_j = sizeJ/NBLOCKS + 1;
	for (int i=threadIdx.x*N_i;i<min(sizeI,(threadIdx.x+1)*N_i);++i) {
		float M1[4];
		float M[4];
		M1[0] = i*step;
		M1[3] = 1;
		for (int j= blockIdx.x*N_j;j<min(sizeJ,(blockIdx.x+1)*N_j);++j) {
			M1[1] = j* step;
			for (int k=1;k<=sizeK;++k) {
				M1[2] = -k*step;
				for (int l=0;l<3;++l) {
					M[l] = 0;
					for (int m=0;m<4;++m) {
						M[l] += smat[l*4 + m] * M1[m];
					}
					}
				int x,y,z;
				if (!interpolate) {//round to nearest neighbor
					x = (int) floor( M[0]+0.5 );
					y = (int) floor( M[1]+0.5 );
					z = (int) floor( M[2]+0.5 );
				}
				else {
					x = (int) floor(M[0]);
					y = (int) floor(M[1]);
					z = (int) floor(M[2]);
				}
				int sizeX = sizeVol2;
				int sizeY = sizeVol1/sizeVol2;
				int sizeZ = sizeVol3/sizeVol1;
				//int index = sizeVol1 * z + sizeVol2 * y + x;
				float value = 0;
				bool isNull = true;
				if (interpolate && x>=0 && x<sizeX-1 && y>=0 && y<sizeY-1 && z>=0 && z<sizeZ-1) {

				//trilinear interpolation, see http://en.wikipedia.org/wiki/Trilinear_interpolation
					float dx,dy,dz;
					dx = (M[0] - x);//:step;
					dy = (M[1] - y);///step;
					dz = (M[2] - z);///step;

					float c000 = (float)tex3D(volumeTexture,x,y,z);					
					float c001 = (float)tex3D(volumeTexture,x,y,z+1);			
					float c010 = (float)tex3D(volumeTexture,x,y+1,z);
					float c011 = (float)tex3D(volumeTexture,x,y+1,z+1);
					float c100 = (float)tex3D(volumeTexture,x+1,y,z);
					float c101 = (float)tex3D(volumeTexture,x+1,y,z+1);
					float c110 = (float)tex3D(volumeTexture,x+1,y+1,z);
					float c111 = (float)tex3D(volumeTexture,x+1,y+1,z+1);
					float i1 = c000 * (1-dz) + c001 * dz;
					float i2 = c010 * (1-dz) + c011 * dz;
					float j1 = c100 * (1-dz) + c101 * dz;
					float j2 = c110 * (1-dz) + c111 * dz;
					float w1 = i1   * (1-dy) + i2   * dy;
					float w2 = j1   * (1-dy) + j2   * dy;
					value    = w1   * (1-dx) + w2   * dx;
				}
				else if (x>=0 && x<sizeX && y>=0 && y<sizeY && z>=0 && z<sizeZ) {
					value = (float)tex3D(volumeTexture,x,y,z);		
				}
				isNull &= (value==0);
				if (!isNull) {
					if (brightness < 0)
						value *= 1 + brightness;
					else
						value += (255-value) * brightness;
					value = (value - 128) * tanf((contrast+1) * PI / 4.) + 128;
				}
				unsigned char value2(value);
				volume_out[sizeJ*sizeK*i+sizeK*j+k] = value2;
			}
		}
	}		
}

__device__ float3 Talairach(float3 in, float3 tp[9]){
	
	float tempx = in.x;
	in.x = in.y;
	in.y = - tempx;
	enum talpts {AC,PC,IHP,AP,PP,SP,IP,RP,LP,EX,EY,EZ};

	float X1 = length(tp[PP] - tp[PC]);
	float X2 = length(tp[AC] - tp[PC]);
	float X3 = length(tp[AP] - tp[AC]);

	float Y1 = length(tp[AC] - tp[RP]);
	float Y2 = length(tp[LP] - tp[AC]);

	float Z1 = length(tp[AC] - tp[IP]);
	float Z2 = length(tp[SP] - tp[AC]);


	float X1T = 79;
	float X2T = 23;
	float X3T = 70;
	float Y1T = 68;
	float Y2T = 68;
	float Z1T = 42;
	float Z2T = 74;
	

	float3 B = in - tp[AC];
	float3 Pac = make_float3(tp[EX].x * B.x + tp[EX].y * B.y + tp[EX].z * B.z,
							 tp[EY].x * B.x + tp[EY].y * B.y + tp[EY].z * B.z,
							 tp[EZ].x * B.x + tp[EZ].y * B.y + tp[EZ].z * B.z);
	//  Transformation from AC PC to Talairach
	float3 Pt = make_float3(0,0,0);
	if (Pac.x < -X2)
		Pt.x = (X1T / X1) * (Pac.x + X2) - X2T;

	if (Pac.x < 0 && Pac.x >= -X2)
		Pt.x = (X2T / X2) * Pac.x;

	if (Pac.x >= 0)
		Pt.x = (X3T / X3) * Pac.x;

	if (Pac.y < 0)
		Pt.y = (Y1T / Y1) * Pac.y;

	if (Pac.y >= 0)
		Pt.y = (Y2T / Y2) * Pac.y;

	if (Pac.z < 0)
		Pt.z = (Z1T / Z1) * Pac.z;

	if (Pac.z >= 0)
		Pt.z = (Z2T / Z2) * Pac.z;
    
	float temp = Pt.y;
	Pt.y = Pt.x;
	Pt.x = temp;
	return Pt;
}


template class IcemVolume<uint8_t>;
template class IcemVolume<uint16_t>;